import tensorflow as tf
import numpy as np

# 设置参数
vocab_size = 10000  # 词汇表大小，和训练时保持一致
maxlen = 500  # 每条评论最大长度，和训练时保持一致

# 加载 IMDB 数据集中的单词索引
word_index = tf.keras.datasets.imdb.get_word_index()

# 添加特殊标记（和训练时一样）
word_index = {k: (v + 3) for k, v in word_index.items()}
word_index["<pad>"] = 0
word_index["<start>"] = 1
word_index["<unk>"] = 2
word_index["<unused>"] = 3

# 反向索引，方便调试
reverse_word_index = dict([(value, key) for (key, value) in word_index.items()])

# 加载模型
model = tf.keras.models.load_model('./models/best_model.keras')
print("✅ 模型加载完成！")


# 处理输入文本
def encode_sentence(sentence):
    words = sentence.lower().split()
    encoded = [word_index.get(w, 2) for w in words]  # 未知词用 2 ('<unk>')

    # 防止索引越界
    encoded = [idx if idx < vocab_size else 2 for idx in encoded]

    return tf.keras.preprocessing.sequence.pad_sequences([encoded], maxlen=maxlen)


# 开始推理
while True:
    text = input("请输入英文评论（输入 exit 退出）：\n")
    if text.strip().lower() == "exit":
        break

    x = encode_sentence(text)
    pred = model.predict(x, verbose=0)[0][0]
    print(f"预测正面情感的概率：{pred:.3f}")
    print("判定：", "正面" if pred > 0.5 else "负面")
